#!/bin/env python3
import argparse
import pickle
import matplotlib.pyplot as plt
import ithildin as ith
import numpy as np
import gc
from pydiffmap.diffusion_map import DiffusionMap
import pydiffmap.kernel as kernel
import matplotlib

parser = argparse.ArgumentParser(description="Make plots about a diffusion map.", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('diffmap', help='diffusion map to load')
parser.add_argument('output_stem')
parser.add_argument('--ith-stem', help='Ithildin stem with points to fit')
parser.add_argument('--fit', type=int, help='fit this number of points')
parser.add_argument('--interactive', action=argparse.BooleanOptionalAction, default=False)
parser.add_argument('--axes', action=argparse.BooleanOptionalAction, help='Keep/remove all decorations (axis names, ticks, margin, border). Does not influence the file name.  Currently only supported for correlation plots and scatter plots.')
subparsers = parser.add_subparsers(required=True,dest='subcommand')
parser_eigenvalues = subparsers.add_parser('eigenvalues', help='plot the eigenvalues')
parser_transformed_scatterplot = subparsers.add_parser('transformed scatterplot', help='plot the transformed coordinates')
parser_correlation_plots = subparsers.add_parser('correlation plots', help='plot the transformed coordinates against the untransformed coordinates')
#parser_scattereplots = subparsers.add_parser('scatterplots', help='plot the untransformed coordinates against each other')

args = parser.parse_args()

print("Reading the diffusion map.")
with open(args.diffmap, 'rb') as f:
    dmap_info = pickle.load(f)
print("Read.")

if not (args.fit is None):
    output_stem = "%s-k%s--etimeshunderd-%s--normalize-%s--N-%s--M-%s--spaceclip-%s--starttime-%s--endtime-%s" % (args.output_stem, dmap_info['local_kernel']['k'],int(100*dmap_info['local_kernel']['epsilon']),dmap_info['normalize'],len(dmap_info['preprocessed_data']),args.fit,dmap_info['space_clip_percentage'],dmap_info['start_time'],dmap_info['end_time'])
else:
    output_stem = "%s-k%s--etimeshunderd-%s--normalize-%s--N-%s--spaceclip-%s--starttime-%s--endtime-%s" % (args.output_stem, dmap_info['local_kernel']['k'],int(100*dmap_info['local_kernel']['epsilon']),dmap_info['normalize'],len(dmap_info['preprocessed_data']),dmap_info['space_clip_percentage'],dmap_info['start_time'],dmap_info['end_time'])

print("Writing to %s..." % (output_stem,))

if args.interactive:
    matplotlib.use("TkAgg")

if args.subcommand == 'eigenvalues':
    fig = plt.figure()
    ax = fig.add_subplot()
    ax.scatter(list(range(1,dmap_info['n_evecs'] + 1)), dmap_info['evals'])
    plt.xticks(list(range(1,dmap_info['n_evecs'] + 1))) # TODO untested
    print(dmap_info['n_evecs'])
    print(dmap_info['evals'])
    fig.savefig("%s-eigenvalues.png" % (output_stem,))
    if args.interactive:
        plt.show()
    plt.close()

    
def apply_percentage(p, x):
    assert(0 <= p)
    assert(p <= 100)
    assert(0 <= x)
    return (p * x)//100 # the exact rounding is expected to only have neglible results

def start_percentage(dimension):
    if dimension==0: # time dimension
        return dmap_info['start_time']
    else: # space dimension
        return dmap_info['space_clip_percentage']//2

def end_percentage(dimension):
    if dimension==0: # time dimension
        return dmap_info['end_time']
    else: # space dimension
        if data.shape[dimension] <= 1:
            return 100
        # TODO: different, use round-upwards division elsewhere
        else:
            return 100 - (dmap_info['space_clip_percentage']//2)

def make_borders(data, start_percentage, end_percentage):
    shape = list(data.shape)
    for i in range(0,len(shape)):
        shape[i] = slice(apply_percentage(start_percentage(i),shape[i]), apply_percentage(end_percentage(i),shape[i]))
    return tuple(shape)

def remove_all_borders(data, variable):
    slices = make_borders(data, start_percentage, end_percentage)
    return variable[slices]

evecs = dmap_info['evecs']
if not (args.fit is None):
    # construct-diffusion-map
    data = ith.SimData.from_stem(args.ith_stem)
    np.random.seed(1)
    vars = dict()
    keys = list(data.vars.keys());
    for key in keys:
        vars[key] = remove_all_borders(data, data.vars[key])
        # save memory
        del data.vars[key]
    # Use the phase space, not the configuration space.
    original_number_of_points = None
    for key in keys:
        vars[key] = vars[key].ravel()
        original_number_of_points = vars[key].shape[0]
    print(original_number_of_points)
    # replace=True necessary due to large number of points
    choices = np.random.choice(original_number_of_points, args.fit, replace=True)
    for key in data.vars.keys():
        vars[key] = vars[key][choices]
    del choices
    gc.collect()
    # normalize
    print('Normalising')
    if dmap_info['normalize']:
        for key in keys:
            if dmap_info['spreads'][key] != 0:
                vars[key] -= dmap_info['means'][key]
                vars[key] /= dmap_info['spreads'][key]
                gc.collect()
    variable_names = list(vars.keys())
    variable_names.sort()
    variable_values = list(variable_names)
    for i in range(0,len(variable_names)):
        variable_values[i] = vars[variable_names[i]]
    variable_values = np.column_stack(tuple(variable_values))
    del vars # save memory
    dmap_reconstructed = DiffusionMap.from_sklearn()
    dmap_reconstructed.dmap = dmap_info['dmap']
    dmap_reconstructed.data = dmap_info['preprocessed_data']
    dmap_reconstructed.right_norm_vec = dmap_info['right_norm_vec']
    dmap_reconstructed.epsilon_fitted = dmap_info['epsilon_fitted']
    dmap_reconstructed.weights = dmap_info['weights']
    dmap_reconstructed.kernel_matrix = dmap_info['kernel_matrix']
    dmap_reconstructed.L = dmap_info['L']
    dmap_reconstructed.q = dmap_info['q']
    dmap_reconstructed.local_kernel = kernel.Kernel()
    dmap_reconstructed.local_kernel.data = dmap_info['preprocessed_data']
    dmap_reconstructed.local_kernel.bandwidth_fxn = None
    dmap_reconstructed.local_kernel.neigh = dmap_info['local_kernel']['neigh']
    dmap_reconstructed.local_kernel.epsilon = dmap_info['local_kernel']['epsilon']
    dmap_reconstructed.local_kernel.epsilon_fitted = dmap_info['local_kernel']['epsilon_fitted']
    dmap_reconstructed.local_kernel.epsilon = dmap_info['local_kernel']['epsilon']
    dmap_reconstructed.local_kernel.k = dmap_info['local_kernel']['k']
    gc.collect()
    print('Transforming')
    evecs = dmap_reconstructed.transform(variable_values)
    print('Transformed!')
    
if args.subcommand == 'transformed scatterplot':
    if dmap_info['n_evecs'] > 3:
        for j in range(0,dmap_info['n_evecs'] - 3):
            evec_x = evecs[:,j]
            evec_y = evecs[:,j+1]
            evec_z = evecs[:,j+2]
            evec_w = evecs[:,j+3]
            fig = plt.figure(figsize=(6,6))
            ax = fig.add_subplot(111,projection='3d')
            ax.scatter(evec_x,evec_y,evec_z,c=evec_w,cmap='viridis')
            # Lets do the title in (Xe)LaTeX?
            ax.set_xlabel(r'$\psi_%s$' % (j + 1))
            ax.set_ylabel(r'$\psi_%s$' % (j + 2))
            ax.set_zlabel(r'$\psi_%s$' % (j + 3))
            ax.axis('tight') # ? copied from pydiffmap
            ax.set_title(r'$\psi_%s$' % (j + 4))
            fig.savefig("%s-transformed-scatterplot--colour%s.png" % (output_stem, j+4))
    else:
        fig = plt.figure(figsize=(6,6))
        ax = fig.add_subplot()
        if dmap_info['n_evecs'] == 2:
            ax.scatter(evecs[:,0],evecs[:,1])
        else:
            ax.scatter(evecs[:,0],evecs[:,1],c=evecs[:,2],cmap='viridis')
            ax.set_title(r'$\psi_3$')
        ax.set_xlabel(r'$\psi_1$')
        ax.set_ylabel(r'$\psi_2$')
        fig.savefig("%s-transformed-scatterplot.png" % (output_stem,))
    if args.interactive:
        plt.show()
    plt.close()

if args.subcommand == 'correlation plots':
    for j in range(0,dmap_info['n_evecs']):
        for i in range(0,len(dmap_info['variable_names'])):
            fig = plt.figure(figsize=(6,6))
            ax = fig.add_subplot()
            var = dmap_info['preprocessed_data'][:,i]
            print("(%s,%s,nmin=%s,nmax=%s,s=%s)" % (i, dmap_info['variable_names'][i], min(var),max(var),dmap_info['spreads'][dmap_info['variable_names'][i]]))
            if dmap_info['normalize']:
                var_min = min(var) * dmap_info['spreads'][dmap_info['variable_names'][i]] + dmap_info['variable_names'][i]
                var_max = max(var) * dmap_info['spreads'][dmap_info['variable_names'][i]] + dmap_info['variable_names'][i]
            else:
                var_min = min(var)
                var_max = max(var)
            # avoid precision issues when not printing axes
            if args.axes and dmap_info['normalize'] and (dmap_info['spreads'][dmap_info['variable_names'][i]] != 0):
                print("normalize-correct")
                var *= dmap_info['spreads'][dmap_info['variable_names'][i]]
                var += dmap_info['means'][dmap_info['variable_names'][i]]
            ax.scatter(dmap_info['dmap'][:,j], dmap_info['preprocessed_data'][:,i])
            print("axes?")
            print(args.axes)
            if args.axes:
                # x and y labels turn out to be redundant for the figures in the thesis, remove them!
                #ax.set_xlabel(r'$\psi_%s$' % (j + 1), fontsize='xx-large')
                #ax.set_ylabel(dmap_info['variable_names'][i], fontsize='xx-large')
                # , "%s" % (var_max,)]
                plt.yticks([min(dmap_info['preprocessed_data'][:,i]), max(dmap_info['preprocessed_data'][:,i])], labels=["%.4f" % (var_min,), "%.4f" % (var_max,)],fontsize='xx-large')
                print("yticks")
            else:
                plt.axis('off')
                ax.xaxis.set_visible(False)
                ax.yaxis.set_visible(False)
                #set_major_locator(matplotlib.ticker.NullLocator())
                #ax.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
                #ax.
            if args.axes:
                plt.xticks([min(dmap_info['dmap'][:,j]), max(dmap_info['dmap'][:,j])], fontsize='xx-large')
            fig.savefig("%s-correlation-plot-%s--eigenvector-%s.png" % (output_stem, dmap_info['variable_names'][i], j + 1), pad_inches = 0, bbox_inches='tight')
            if args.interactive:
                plt.show()
            plt.close()
            
